// implement fork from user space

#include <inc/string.h>
#include <inc/lib.h>

// PTE_COW marks copy-on-write page table entries.
// It is one of the bits explicitly allocated to user processes (PTE_AVAIL).
#define PTE_COW		0x800

extern void _pgfault_upcall(void);

// 通过UVPT查询va对应的pte, 如果没有返回null
pte_t*
get_pte_for_va(void* va) {
    pte_t* pte = 0;

    pde_t* pde = (pde_t*)(PGADDR(PDX(UVPT), PDX(UVPT), 4*PDX(va)));
    if (*pde & PTE_P) {
        // page table exists
        pte = (pte_t*)(PGADDR(PDX(UVPT), PDX(va), 4*PTX(va)));
    }

    return pte;
}

pde_t*
get_pde_for_va(void* va) {
    return (pde_t*)(PGADDR(PDX(UVPT), PDX(UVPT), 4*PDX(va)));
}
//
// Custom page fault handler - if faulting page is copy-on-write,
// map in our own private writable copy.
//
static void
pgfault(struct UTrapframe *utf)
{
	void *addr = (void *) utf->utf_fault_va;
	uint32_t err = utf->utf_err;
	int r;
    //cprintf("#PF @ [0x%x]\n", addr);

	// Check that the faulting access was (1) a write, and (2) to a
	// copy-on-write page.  If not, panic.
	// Hint:
	//   Use the read-only page table mappings at uvpt
	//   (see <inc/memlayout.h>).

	// LAB 4: Your code here.
    // check write operation
    if (!(err & FEC_WR)) {
        cprintf("#PF at [0x%x] by va:%x\n", utf->utf_eip, (uintptr_t)addr);
        panic("page fault not triggered by a write" );
    }

    // check page
    /*
     * 如何通过uvpt查看指定页的pte?
     *
     * 先来看一下给定一个va, 如何找到va对应的物理地址
     * 1. 通过cr3找到page dir的物理地址
     * 2. va的高10bit作为偏移量, 在page dir 中找到对应的pde, 里面有
     *    page table的物理地址
     * 3. va的次高10bit作为偏移量, 在page table中找到对应的pte, 里面有va所属
     *    的物理页框
     * 4. va中的低12bit作为偏移量, 结合找到的物理页框, 可以确定va物理地址
     *
     * 再来看通过uvpt找到页面的pte, 假设需要找va的pte, va分为如下部分
     *   page dir  |  page table  |  offset
     *
     * 可以通过构造如下的qva(query va)来获取va的pte
     *   PDX(UVPT) |  page dir    |  page table * 4
     *
     * 结合下面链接中的图, 很容易验证
     * https://pdos.csail.mit.edu/6.828/2018/labs/lab4/uvpt.html
     */

    pte_t* fault_page_pte = get_pte_for_va(addr);
    if (!fault_page_pte)
        panic("no such page");
    if (!(*fault_page_pte & PTE_COW))
        panic("page fault not triggered by COW page");


	// Allocate a new page, map it at a temporary location (PFTEMP),
	// copy the data from the old page to the new page, then move the new
	// page to the old page's address.
	// Hint:
	//   You should make three system calls.

	// LAB 4: Your code here.
    envid_t envid = sys_getenvid();
    if ((r = sys_page_alloc(envid, PFTEMP, PTE_W|PTE_U|PTE_P) < 0)){
        cprintf("error: %e\n", r);
        panic("error allocating page");
    }
    memcpy(PFTEMP, (void*)(PGNUM(addr)<<PTXSHIFT) , PGSIZE);
    if ((r = sys_page_map(envid, PFTEMP, envid,
            (void*)(PGNUM(addr)<<PTXSHIFT), 
            PTE_W|PTE_U|PTE_P)) < 0){
        cprintf("error: %e\n", r);
        panic("error map page");
    }
    sys_page_unmap(envid, PFTEMP);

}

//
// Map our virtual page pn (address pn*PGSIZE) into the target envid
// at the same virtual address.  If the page is writable or copy-on-write,
// the new mapping must be created copy-on-write, and then our mapping must be
// marked copy-on-write as well.  (Exercise: Why do we need to mark ours
// copy-on-write again if it was already copy-on-write at the beginning of
// this function?)
//
// Returns: 0 on success, < 0 on error.
// It is also OK to panic on error.
//
static int
duppage(envid_t envid, unsigned pn)
{
	int r;

	// LAB 4: Your code here.
    pte_t* pte = get_pte_for_va((void*)(pn*PGSIZE));
    if (!pte)
        return 0;

    envid_t thisenvid = sys_getenvid();
    int perm = PTE_U | PTE_P;
    // lab5
    if ((*pte & PTE_SHARE)) {
        // 原样复制, 这是共享的页
        //cprintf("copy shared page [0x%x]\n", pn*PGSIZE);
        sys_page_map(thisenvid, (void*)(pn*PGSIZE), envid,
                (void*)(pn*PGSIZE), perm | PTE_SHARE | PTE_W);
    }
    else if ((*pte & PTE_COW) || (*pte & PTE_W) ) {
        perm |= PTE_COW;
        r = sys_page_map(thisenvid, (void*)(pn*PGSIZE),
                envid, (void*)(pn*PGSIZE), perm);
        if (r < 0) {
            cprintf("%e\n", r);
            panic("error duppage, cannot map thisenv -> envid");
        }

        // mark this env as pow
        sys_page_map(thisenvid, (void*)(pn*PGSIZE),
                thisenvid, (void*)(pn*PGSIZE), perm);
        if (r < 0)
            panic("error duppage, cannot mark thisenv as pow");

    } else {
        if (sys_page_map(thisenvid, (void*)(pn*PGSIZE),
                envid, (void*)(pn*PGSIZE), perm) < 0)
            panic("error duppage, cannot map non-pow page");
    }


	return 0;
}

//
// User-level fork with copy-on-write.
// Set up our page fault handler appropriately.
// Create a child.
// Copy our address space and page fault handler setup to the child.
// Then mark the child as runnable and return.
//
// Returns: child's envid to the parent, 0 to the child, < 0 on error.
// It is also OK to panic on error.
//
// Hint:
//   Use uvpd, uvpt, and duppage.
//   Remember to fix "thisenv" in the child process.
//   Neither user exception stack should ever be marked copy-on-write,
//   so you must allocate a new page for the child's user exception stack.
//
envid_t
fork(void)
{
    int r;
	// LAB 4: Your code here.

    // fork分如下几步
    // 1. 注册pgfault处理函数
    // 2. 创建出子进程
    // 3. 复制父进程的页表到子进程, 标记COW
    // 4. 为子进程注册pgfault函数
    // 5. 标记子进程为runnable

    // 1. 注册pgfault
    set_pgfault_handler(pgfault);

    // 2. exofork
    envid_t envid;
    if ((envid = sys_exofork()) < 0)
        // error
        panic("error forking");

    if (envid == 0) {
        // child
        // 修改thisenv, 不然过不了stresssched测试的
        thisenv = &envs[ENVX(sys_getenvid())];
    } else {
        // parent
        // 3. 复制页表
        for (int i = 0 ; i < UXSTACKTOP/PGSIZE - 1 ;) { 
            // which I will alloc manually
            pde_t* pde = get_pde_for_va((void*) (i*PGSIZE));
            if (!(*pde & PTE_P) || !(*pde & PTE_U)) {
                // pde不存在, 跳过1024个页
                i += PGSIZE/sizeof(pte_t);
                i &= ~(PGSIZE/sizeof(pte_t)-1);
            } else {
                pte_t* pte = get_pte_for_va((void*) (i*PGSIZE));
                if (pte && (*pte & PTE_P) && (*pte & PTE_U)) {
                    Debug("dup page[0x%x]\n", i*PGSIZE);
                    duppage(envid, i);
                }

                i++;
            }
        }

        Debug("allocating exception stack\n");
        // user exception stack
        if ((r = sys_page_alloc(envid, (void*)(UXSTACKTOP-PGSIZE), 
                        PTE_P|PTE_W|PTE_U)) < 0) {
            cprintf("error allocate exception stack for env #%x:%e\n",
                    envid, r);
            panic("error allocate");
        }

        // 4. 注册pgfault
        Debug("registering pgfault for child\n");
        sys_env_set_pgfault_upcall(envid, _pgfault_upcall);

        // 5. 标记为可执行
        Debug("mark child as runnable\n");
        sys_env_set_status(envid, ENV_RUNNABLE);
    }
    return envid;
}

// Challenge!
int
sfork(void)
{
	panic("sfork not implemented");
	return -E_INVAL;
}
